package easy

import (
	"errors"
	"net"
	"strings"
	"sync"
	"time"

	"gitee.com/haodreams/libs/routine"
)

type BlackHost struct {
	Host   string `json:"host"`
	Reason string `json:"reason"`
	Time   int64  `json:"time"`
}

// 限制ip访问
type BlackListener struct {
	net.Listener
	blackMap map[string]*BlackHost
	lock     sync.RWMutex
	logf     func(v ...any)
	timeout  int64
}

func (m *BlackListener) IsBlackHost(host string) bool {
	m.lock.RLock()
	hi := m.blackMap[host]
	m.lock.RUnlock()
	if hi != nil {
		if time.Now().Unix()-hi.Time > m.timeout {
			m.DeleteBlack(hi.Host)
			return false
		}
		return true
	}
	return false
}

func NewBlackListener(logf func(v ...any), timeouts ...int64) *BlackListener {
	m := new(BlackListener)
	m.logf = logf
	m.blackMap = make(map[string]*BlackHost)
	m.lock = sync.RWMutex{}
	m.timeout = 10 * 60
	if len(timeouts) > 0 {
		m.timeout = timeouts[0]
	}
	if m.timeout < 15 {
		m.timeout = 15
	}
	return m
}

func (m *BlackListener) Accept() (conn net.Conn, err error) {
	for routine.IsRunning() {
		conn, err = m.Listener.Accept()
		if err != nil {
			return
		}

		// 获取请求的 IP 地址
		clientIP := conn.RemoteAddr().String()
		pos := strings.LastIndex(clientIP, ":")
		if pos < 0 {
			continue
		}
		clientIP = clientIP[:pos]
		if clientIP == "127.0.0.1" {
			return
		}

		host := m.GetBlack(clientIP)
		if host == nil {
			return
		}
		if time.Now().Unix()-host.Time > m.timeout {
			m.DeleteBlack(host.Host)
			return
		}

		if m.logf != nil {
			m.logf("禁止访问:", clientIP)
		}
		conn.Close()
	}
	return
}

func (m *BlackListener) AcceptTCP() (conn *net.TCPConn, err error) {
	c, err := m.Accept()
	if err != nil {
		return
	}
	conn, ok := c.(*net.TCPConn)
	if !ok {
		err = errors.New("类型错误")
		return
	}
	return
}

func (m *BlackListener) AddBlackHost(host string, reason string) {
	m.lock.RLock()
	hi := m.blackMap[host]
	m.lock.RUnlock()
	if hi == nil {
		hi = new(BlackHost)
		hi.Host = host
	}
	hi.Reason = reason
	hi.Time = time.Now().Unix()
	m.lock.Lock()
	m.blackMap[host] = hi
	m.lock.Unlock()
}

func (m *BlackListener) GetBlack(host string) (hi *BlackHost) {
	m.lock.RLock()
	hi = m.blackMap[host]
	m.lock.RUnlock()
	return
}

func (m *BlackListener) DeleteBlack(host string) (hi *BlackHost) {
	m.lock.Lock()
	delete(m.blackMap, host)
	m.lock.Unlock()
	return
}
